# core/huatuogpt_vision.py
import os
import sys
import tempfile
from pathlib import Path
from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import JSONResponse
import uvicorn

# 核心路径配置 ------------------------------------------------------------ 
_current_file = Path(__file__).resolve()
_project_root = _current_file.parent.parent  # 项目根目录
sys.path.insert(0, str(_project_root))

# 模型服务类 -------------------------------------------------------------- 
class HuatuoVisionService:
    def __init__(self):
        """同步初始化所有资源"""
        self.model_dir = _project_root / "core" / "cli" / "huatuogpt-vision-7b"
        self.device = "cuda:0"
        self._verify_model_files()
        self._init_model()
        print("[服务初始化] 华佗视觉模型已加载")

    def _verify_model_files(self):
        """严格模型文件验证"""
        required_files = {
            "config.json": "模型配置文件",
            "model.safetensors.index.json": "模型索引文件",
            "vit/config.json": "视觉编码器配置"
        }
        
        for file, desc in required_files.items():
            path = self.model_dir / file
            if not path.exists():
                raise FileNotFoundError(f"{desc}缺失: {path}")
        print("[文件验证] 所有必需模型文件已确认")

    def _init_model(self):
        """同步初始化模型实例"""
        from core.cli.cli import HuatuoChatbot  # 延迟导入
        print("[模型加载] 正在加载模型权重...")
        self.bot = HuatuoChatbot(
            model_dir=str(self.model_dir),
            device=self.device
        )
        # 执行健康检查
        # test_response = self.bot.chat("健康检查", [])
        # print(f"[健康测试] 模型响应样例: {test_response[:50]}...")

# 全局服务实例 ------------------------------------------------------------ 
try:
    service = HuatuoVisionService()  # 同步初始化
except Exception as e:
    print(f"[致命错误] 服务初始化失败: {str(e)}")
    sys.exit(1)

# FastAPI应用配置 -------------------------------------------------------- 
app = FastAPI(title="华佗视觉医疗助手")

# 用于存储临时文件的列表
temp_files = []

# API端点 --------------------------------------------------------------- 
@app.post("/api/chat")
async def chat_endpoint(
    text: str = Form(..., min_length=1, max_length=1000),
    images: list[UploadFile] = File(...)
):
    """增强型医疗问答接口"""
    # 输入验证
    if len(images) > 6:
        raise HTTPException(400, "最多支持6张医疗影像")
    
    # 处理临时文件
    try:
        for img in images:
            # 文件类型验证
            if img.content_type not in ["image/jpeg", "image/png"]:
                raise HTTPException(400, "仅支持JPEG/PNG格式影像")
            
            # 创建临时文件
            suffix = Path(img.filename).suffix
            with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as f:
                content = await img.read()
                f.write(content)
                temp_files.append(f.name)  # 将临时文件路径添加到列表
                print(f"[临时文件] 已保存: {f.name}")

        # 执行同步推理
        response = service.bot.chat(text, temp_files)
        return JSONResponse({
            "response": response,
            "status": "success"
        })
        
    except Exception as e:
        print(f"[错误] 推理失败: {str(e)}")
        raise HTTPException(500, "内部服务器错误")

@app.get("/health")
def health_check():
    """健康检查接口"""
    return {
        "status": "running",
        "model_loaded": service.bot is not None,
        "model_path": str(service.model_dir),
        "max_images": 6
    }

@app.on_event("shutdown")
def cleanup_temp_files():
    """在API关闭时清理临时文件"""
    print("[清理临时文件] 开始清理...")
    for path in temp_files:
        try:
            os.unlink(path)
            print(f"[清理临时文件] 已删除: {path}")
        except Exception as e:
            print(f"[警告] 文件清理失败: {path} - {str(e)}")
    temp_files.clear()  # 清空临时文件列表
    print("[清理临时文件] 完成")

if __name__ == "__main__":
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=8000,
        # 生产环境建议添加以下参数
        # workers=1,  # 禁用多worker模式
        # timeout_keep_alive=30
    )